from llmtcg.agents.base import Agent
from llmtcg.chat.model import ChatContext, Role

from loguru import logger

from dataclasses import dataclass, field
from typing import List, Tuple, Optional
from tempfile import mkstemp, mkdtemp
from subprocess import run, PIPE, CompletedProcess, STDOUT, TimeoutExpired
from os import path, makedirs

from shutil import rmtree


CCODER_SYSTEM_PROMPT = """
You are a professional programmer and are supposed to write a program to
achieve the given requirement. Your response should start with a brief
introduction of how the program works, and then the generated program.
The generated program should be surrounded by "```c" before it and
"```" after it.

You should always carefully check whether the header files are properly included
if you have used some function that are declared in them. For example, if you have
used the function "strlen", you need to add "#include <string.h>" at the beginning.
"""

CCODER_RETRY_PROMPT = """
Your generated program encounters the following problem: {problem_description},
detailed message is shown as follows.

```
{feedback}
```

Please fix the previous program and re-generate the correct one that matches the original
requirement.
"""

CCODER_FAKE_RESP = """
The program is generated as follows,

```c
{program}
```
"""


class CodingFailed(Exception):
    pass


@dataclass
class CCoderAgent(Agent):

    compiler: str = "gcc"
    compiler_args: List[str] = field(default_factory=list)
    max_retries: int = 5
    with_history: bool = True

    _last_context: Optional[ChatContext] = None

    @property
    def system_prompt(self) -> str:
        return CCODER_SYSTEM_PROMPT

    @property
    def last_context(self) -> ChatContext:
        assert self._last_context is not None, f"no last context found"
        return self._last_context

    def validate_program(self, program: str, tfname: str) -> Tuple[bool, bool, str]:
        """
        We use a inherited temp file name tfname from outside world, so that every time
        we compile/run the program, the produced error message has the same source file
        name.
        """
        with open(tfname, "w+") as tf:
            tf.write(program)

        outname = f"{tfname}.out"
        cproc: CompletedProcess = run(
            [self.compiler, *self.compiler_args, tfname, "-o", outname],
            stdin=PIPE, stdout=PIPE, stderr=PIPE
        )
        if cproc.returncode != 0:
            return False, False, cproc.stderr.decode()

        temp_dir = mkdtemp(prefix="llmtcg-coder-")
        makedirs(temp_dir, exist_ok=True)

        try:
            rproc = run(
                [outname], cwd=temp_dir,
                stdin=PIPE, stdout=PIPE, stderr=STDOUT,
                start_new_session=True, timeout=5
            )
            rmtree(temp_dir, ignore_errors=True)
            if rproc.returncode != 0:
                msg = rproc.stdout.decode()
                return True, False, msg
        except TimeoutExpired:
            msg = """ Running this program lead to a timeout problem, probably because
            the program is waiting for something from another process or stdin that will
            never come.
            """
            rmtree(temp_dir, ignore_errors=True)
            return True, False, msg
        except UnicodeDecodeError:
            rmtree(temp_dir, ignore_errors=True)
            return True, False, "failed to decode the output of the test case"

        return True, True, ""

    def get_program(self, requirement: str) -> str:
        if self.with_history:
            return self.__get_program_with_history(requirement)
        else:
            return self.__get_program_without_history(requirement)

    def __extract_program(self, response: str) -> str:
        try:
            start_pos = response.index("```c")
            end_pos = response.index('```', start_pos + 1)
        except ValueError:
            raise CodingFailed("cannot extract prgoram from response")

        program = response[start_pos + 4: end_pos].strip()
        return program

    def __get_program_with_history(self, requirement: str) -> str:
        _, tfname = mkstemp(prefix="llmtcg-agent-ccoder", suffix=".c")

        chat_ctx = ChatContext(self.chat_model)
        chat_ctx.task_identifier = "CCoderAgent::getProgram"
        chat_ctx.system_prompt = self.system_prompt

        self._last_context = chat_ctx

        try:
            resp = chat_ctx.chat(f"Please write a c program according to the given requirement: {requirement}")
        except Exception as ex:
            raise CodingFailed(f"chat API failed because: {ex}")

        nretry = 0
        last_error = ""

        while True:
            try:
                start_pos = resp.index("```c")
                end_pos = resp.index('```', start_pos + 1)
            except ValueError:
                logger.debug(resp)
                raise CodingFailed("cannot extract prgoram from response")

            program = resp[start_pos + 4: end_pos].strip()
            compilable, runnable, msg = self.validate_program(program, tfname)

            if compilable and runnable:
                """ We are good ;) """
                logger.info(f"c program generated with {nretry} retries")
                return program

            if msg == last_error:
                raise CodingFailed(f"the exact error message reported again, quit")

            """ Make a retry ;) """
            nretry += 1

            if nretry > self.max_retries:
                raise CodingFailed(f"retry exceed the limitation")

            logger.debug('------ old program ------')
            logger.debug(program)

            if not compilable:
                logger.debug('----- compiling error -----')
                logger.debug(msg.strip())
                new_prompt = f"""
                Please correct the previous program you wrote, it cannot be compiled.
                {self.compiler} reported the following error:

                    {msg.strip()}
                """
            else:
                logger.debug('----- running error -----')
                logger.debug(msg.strip())
                new_prompt = f"""
                Please correct the previous program you wrote, it can be successfully
                compiled, but running it lead to the following error:

                    {msg.strip()}
                """

            try:
                resp = chat_ctx.chat(new_prompt)
            except Exception as ex:
                raise CodingFailed(f"chat API failed because: {ex}")

    def __get_program_without_history(self, requirement: str) -> str:
        _, tfname = mkstemp(prefix="llmtcg-agent-ccoder", suffix=".c")

        chat_ctx = ChatContext(self.chat_model)
        chat_ctx.task_identifier = "CCoderAgent::getProgram"
        chat_ctx.system_prompt = self.system_prompt

        self._last_context = chat_ctx

        initial_prompt = f"Please write a c program according to the given requirement: {requirement}"
        try:
            resp = chat_ctx.chat(initial_prompt)
        except Exception as ex:
            raise CodingFailed(f"chat API failed because: {ex}")

        nretry = 0
        last_error = ""

        while True:
            program = self.__extract_program(resp)
            compilable, runnable, msg = self.validate_program(program, tfname)

            if compilable and runnable:
                """ We are good ;) """
                logger.info(f"c program generated with {nretry} retries")
                return program

            if msg == last_error:
                raise CodingFailed(f"the exact error message reported again, quit")

            """ Make a retry ;) """
            nretry += 1

            if nretry > self.max_retries:
                raise CodingFailed(f"retry exceed the limitation")

            if not compilable:
                problem_description = f"the program cannot be compiled, {self.compiler} reported the following error"
            else:
                problem_description = \
                    "the program can be successfully compiled, " + \
                    "but running it lead to the following error"

            retry_prompt = CCODER_RETRY_PROMPT.format(
                problem_description=problem_description,
                feedback=msg
            )

            try:
                chat_ctx.clear()
                chat_ctx.history.append((Role.USER, initial_prompt))
                chat_ctx.history.append((Role.ASSISTANT, CCODER_FAKE_RESP.format(program=program)))
                resp = chat_ctx.chat(retry_prompt)
            except Exception as ex:
                raise CodingFailed(f"chat API failed because: {ex}")
